🚩 Coordinator startup ordering relies on JAX retry semantics

Workers (rank > 0) start their JAX backend process immediately at modal_train.py:160, potentially before rank 0 has even started the API server / JAX coordinator. The JAX backend workers connect to --coordinator-address and must internally retry until the coordinator is reachable. If the JAX backend does not implement connection retries, workers would fail immediately. This appears to work correctly with SkyRL-TX's JAX backend (which uses JAX's distributed initialization that retries coordinator connections), but it's worth verifying if the SkyRL commit is ever updated.

Was this helpful? React with 👍 or 👎 to provide feedback.

Confirmed this was a concrete ordering robustness issue. I pushed b2c3abc, which gates rank>0 worker startup on a Modal Dict signal written only after rank 0 has spawned the coordinator/API process, so workers no longer launch before the coordinator process exists.

This is outdated on the current branch: _run_worker() now calls _wait_for_coordinator_launch() before spawning the JAX backend process, and rank 0 only writes run_state[run_key] = "running" after the API/coordinator subprocess has been started. So the current code no longer depends on JAX retry semantics for the rank>0 startup ordering.

Addressed in the current branch. Rank >0 now calls _wait_for_coordinator_launch() before spawning skyrl.backends.jax, and rank 0 sets run_state[run_key] = "running" only after starting the API/coordinator process. I also added a SkyRL RPC barrier patch so workers and coordinator enter subsequent multihost calls in the same order.

Update: I removed the experimental RPC barrier patch while narrowing the checkpoint fix. The concrete startup-ordering issue in this thread remains addressed by the run_state gate: rank >0 waits for rank 0 to mark the coordinator process as running before it launches skyrl.backends.jax.

🚩 _commit_and_print_checkpoints depends on .tar.gz format assumption

At modal_train.py:286, _commit_and_print_checkpoints searches for *.tar.gz files under the checkpoint root. If the SkyRL-TX Tinker API server saves checkpoints in a different format (e.g., directory-based Orbax checkpoints or SafeTensors), this glob would find nothing and raise a RuntimeError. The _patch_skyrl_checkpointing function patches the backend to use Orbax, whose native format is directory-based, not .tar.gz. However, the Tinker SDK's save_state and save_weights_for_sampler methods likely wrap this into their own archive format. The PR's README shows successful results with .tar.gz checkpoint files, so this presumably works, but the coupling to a specific archive format is fragile.

Was this helpful? React with 👍 or 👎 to provide feedback.

Acknowledged. This is intentionally scoped to the pinned SkyRL commit (830f005...), whose Tinker checkpoint writer produced .tar.gz artifacts in both validated Modal runs. If the SkyRL pin is updated and the checkpoint format changes, this helper should be updated with that pin bump.

Fixed in c27f75b: _commit_and_print_checkpoints now uses _checkpoint_artifacts(...), which prefers the validated .tar.gz archive outputs but falls back to any non-results checkpoint files under the run-scoped checkpoint root. The log line is now *_checkpoint_artifact=... so future directory/file checkpoint formats won't false-fail just because they are not .tar.gz archives.

🚩 Worker checkpoint_volume.commit() after rank 0 has already committed

Workers call checkpoint_volume.commit() at skyrl-tx/modal_train.py:236 after seeing rank 0's "done" signal. Rank 0 already committed the volume at line 290 (inside _commit_and_print_checkpoints). With the patched Orbax checkpointer using primary_host=None, all JAX processes participate in writing checkpoints. If workers write their portion of a multi-process checkpoint to the volume mount, those writes must also be committed from the worker containers. The worker commit after rank 0's commit ensures this. However, there may be a brief window where the volume is in a partially-committed state between rank 0's commit and workers' commits.

Was this helpful? React with 👍 or 👎 to provide feedback.

I checked this path against the current control flow and the validated Modal runs. _run_client(...) only returns after the server-side save calls finish, so the worker-side checkpoint writes have completed before rank 0 enters _commit_and_print_checkpoints(...). Rank 0 then commits what it can see, sets run_state[run_key] = "done" in finally, and the worker containers each commit their own mounted-volume writes before returning. The clustered Modal call does not complete until those worker containers return, so the post-run durable volume state includes the worker commits.

I agree there can be a transient live-log window between rank 0’s commit log line and the worker commit calls, but that window is before the clustered job has completed and before a follow-up run should consume the volume. The final SFT/RL validation runs both restored/evaluated sampler checkpoints after save, so I’m treating this as a maintenance note rather than a current functional bug.